import base64
import os
import sys
import glob
from pathlib import Path

def encode_file_to_base64(file_path):
    """将文件编码为base64字符串"""
    try:
        with open(file_path, "rb") as f:
            return base64.b64encode(f.read()).decode('utf-8')
    except Exception as e:
        print(f"处理文件 {file_path} 时出错: {e}")
        return None

def generate_function_name(file_name):
    """根据文件名生成函数名"""
    # 移除扩展名并替换非法字符
    base_name = Path(file_name).stem
    # 替换非字母数字字符为下划线
    return f"get_{base_name}_data"

def generate_variable_name(file_name):
    """根据文件名生成变量名，确保以字母开头"""
    # 移除扩展名
    base_name = Path(file_name).stem
    # 添加前缀确保变量名以字母开头
    return f"_{base_name.upper()}_DATA"

def main():
    # 设置输出文件路径
    output_file = "./icon_data.py"
    
    # 获取命令行参数或使用默认图片
    if len(sys.argv) > 1:
        # 用户指定了文件
        image_files = sys.argv[1:]
    else:
        # 提示用户输入文件路径或使用通配符
        print("请输入图片文件路径，多个文件用空格分隔，或使用通配符(如 *.svg *.png):")
        input_paths = input().strip().split()
        
        # 处理通配符
        image_files = []
        for path in input_paths:
            # 展开通配符
            matched_files = glob.glob(path)
            if matched_files:
                image_files.extend(matched_files)
            else:
                # 如果没有匹配到文件，可能是直接的文件路径
                if os.path.isfile(path):
                    image_files.append(path)
    
    if not image_files:
        print("未指定任何有效的图片文件！")
        return
    
    # 开始处理文件
    print(f"将处理以下文件:")
    for file in image_files:
        print(f" - {file}")
    
    # 创建输出文件内容
    output_content = [
        "# 自动生成的图标数据文件",
        "# 包含多个图片的base64编码",
        ""
    ]
    
    # 处理每个文件
    function_definitions = []
    data_variables = []
    
    for file_path in image_files:
        if not os.path.isfile(file_path):
            print(f"警告: {file_path} 不是有效文件，已跳过")
            continue
        
        # 获取文件名和函数名
        file_name = os.path.basename(file_path)
        function_name = generate_function_name(file_name)
        variable_name = generate_variable_name(file_name)
        
        # 编码文件
        base64_data = encode_file_to_base64(file_path)
        if base64_data is None:
            continue
        
        # 添加数据变量
        data_variables.append(f"{variable_name} = \"\"\"\n{base64_data}\n\"\"\"")
        
        # 添加获取函数
        function_definitions.append(
            f"def {function_name}():\n"
            f"    \"\"\"获取 {file_name} 的base64编码数据\"\"\"\n"
            f"    return {variable_name}"
        )
    
    # 组合所有内容
    output_content.extend(data_variables)
    output_content.append("")  # 空行分隔
    output_content.extend(function_definitions)
    
    # 移除旧版兼容函数部分
    
    # 写入文件
    with open(output_file, "w", encoding="utf-8") as f:
        f.write("\n".join(output_content))
    
    print(f"已成功将 {len(function_definitions)} 个图片编码写入到 {output_file}")
    print(f"可用的函数:")
    for func in function_definitions:
        print(f" - {func.split('(')[0]}")

if __name__ == "__main__":
    # 使用方式python generate_base64 文件1 文件2
    main()